from utils.utils import unpickle,pickle
import torch
from models.memory import Memory
import os.path as osp
from scipy.io import loadmat
import json
import numpy as np
from train_unsup import *

def to_device(images, targets, device):
    images = [image.to(device) for image in images]
    for t in targets:
        t["boxes"] = t["boxes"].to(device)
        t["labels"] = t["labels"].to(device)
    return images, targets

#------------------------------------unsup------------------------------------
def create_new_pid_ssm_gt_no_model(cfg,dataset):
    annotations=dataset.annotations
    '''
    annotations.append(
                {
                    "img_name": img_name,
                    "img_path": osp.join(self.img_prefix, img_name),
                    "boxes": boxes,
                    "pids": pids,
                }
            )'''
     # 提取boxes-anno
    pid = 1
    
    num_classes=0
    for i,ann_dict in enumerate(annotations):
        boxes_num = len(ann_dict['boxes'])
        num_classes += boxes_num
        pids = [pid + x for x in range(boxes_num)]  # 3
        pid = pids[-1] + 1  # 4
        annotations[i]['pids']=pids

    mem = torch.zeros([num_classes, 256]) # 55272
    print('pid: ',pid,'loading zeros mem finished..')
    mem = mem.to(cfg.DEVICE)

    num_features = 256
    memory = Memory(num_features, num_classes, mem)
    memory.to(cfg.DEVICE)
    return annotations, memory


#------------------------------------unsup load Lslc model ------------------------------------
def create_new_pid_ssm_gt(cfg,dataset,model=None):
    annotations=dataset.annotations
    '''
    annotations.append(
                {
                    "img_name": img_name,
                    "img_path": osp.join(self.img_prefix, img_name),
                    "boxes": boxes,
                    "pids": pids,
                }
            )'''
     # 提取boxes-anno
    pid = 1
    
    num_classes=0
    for i,ann_dict in enumerate(annotations):
        boxes_num = len(ann_dict['boxes'])
        num_classes += boxes_num
        pids = [pid + x for x in range(boxes_num)]  # 3
        pid = pids[-1] + 1  # 4
        annotations[i]['pids']=pids

    #用初始化的model进行提取 mem
    if model is None:
        mem = torch.zeros([num_classes, 256]) # 55272
        print('pid: ',pid,'loading zeros mem finished..')
        mem = mem.to(cfg.DEVICE)

    else:
        #use cache
        extr_f = '../two_path_init_Lslc_5prw_.pth'
        if osp.exists(extr_f):
            mem = torch.load(extr_f)
            mem = mem.to(cfg.DEVICE)
        else:
            model.eval()
            mem = []
            train_gal_loader = build_train_loader_mlc(cfg, annotations, bs=1)
            for images, targets in tqdm(train_gal_loader, ncols=0):
                images, targets = to_device(images, targets, cfg.DEVICE)
                embeddings = model(images, targets)
                feats = torch.cat(embeddings).cpu().detach()
                mem.append(feats)

            mem = torch.cat(mem)
            mem = mem.to(cfg.DEVICE)
            torch.save(mem, extr_f)

    num_features = 256
    memory = Memory(num_features, num_classes, mem)
    memory.to(cfg.DEVICE)
    return annotations, memory


#ssm-trans-from-prw yoloboxes for anno and prw-reid extract for mem=False
def create_new_pid_ssm_trans_from_prw_yolo(model,cfg):
    model.eval()
    load=cfg.LOAD_MEM

    annotations = []
    mem = []

    # get annotations
    dir='/home/cv7609/zjh/000seqgit_v1/0yolov5_scores/'
    yolobox_json =dir+ cfg.score_anno

    with open(osp.join(yolobox_json), "r") as f:
        yolodet_boxes = json.load(f)  # {'person':[[x1,y1,x2,y2,s,img_path]]}
    # yolobox =yolodet_boxes['person'] #{'person':[[x1,y1,x2,y2,s,img_path]]}
    img_has_boxes = {}  # {path:boxes}
    for box_info in yolodet_boxes['person']:
        box, img_path = box_info[:4], box_info[-1]
        img_has_boxes[img_path] = []
    for box_info in yolodet_boxes['person']:
        box, img_path = box_info[:4], box_info[-1]
        img_has_boxes[img_path].append(box)

    num_classes = len(yolodet_boxes['person'])
    print('yolo 总box数： ', len(yolodet_boxes['person']))

    #重命名single class id box anno
    pid = 1
    for img_p, boxes in img_has_boxes.items():
        img_name = img_p.split('/')[-1]
        pids = [pid + x for x in range(len(boxes))]  # 3
        pid += len(boxes)  # 4
        annotations.append(
            {
                "img_name": img_name,
                "img_path": osp.join(cfg.INPUT.DATA_ROOT, 'Image', 'SSM', img_name),
                "boxes": boxes,
                "pids": pids
            }
        )

    # use cache
    extr_f = '../ssm_feats/extract_ssm_feats_n2w_bs3.pth'
    if osp.exists(extr_f) and load:
        mem = torch.load(extr_f)
        mem = mem.to(cfg.DEVICE)

    # not use cache extract prw-reid mem
    else:
        train_gal_loader = build_train_loader_mlc(cfg, annotations, bs=1)
        for images, targets in tqdm(train_gal_loader, ncols=0):
            images, targets = to_device(images, targets, cfg.DEVICE)
            embeddings = model(images, targets)
            feats = torch.cat(embeddings).cpu().detach()
            mem.append(feats)

        mem = torch.cat(mem)
        mem = mem.to(cfg.DEVICE)
        torch.save(mem, extr_f)

    num_features = 256
    memory = Memory(num_features, num_classes, mem)
    memory.to(cfg.DEVICE)

    return annotations,memory

#prw-unsup use yolobox mem0=True
def create_new_pid_prw(cfg):
    '''
    prw-unsup use yolobox mem0=True
    '''
    # ------------------------------预测prw-train的 dets和feats得到单标签，形成annotation---------------------------------------------------
    data_imgroot =osp.join(cfg.INPUT.DATA_ROOT,'frames')  # '../../ps_raw/prw/frames/'
    dataroot = cfg.INPUT.DATA_ROOT
    imgs = loadmat(osp.join(dataroot, "frame_train.mat"))["img_index_train"]
    img_path = [data_imgroot + '/'+img[0][0] + ".jpg" for img in imgs]

    wid = 0
    nid = 0
    # 对GT box的统计
    for im_p in img_path:
        img_name = im_p.split('/')[-1]
        anno_path = osp.join(dataroot, "annotations", img_name)
        anno = loadmat(anno_path)
        box_key = "box_new"
        if box_key not in anno.keys():
            box_key = "anno_file"
        if box_key not in anno.keys():
            box_key = "anno_previous"

        rois = anno[box_key][:, 1:]
        ids = anno[box_key][:, 0]
        rois = np.clip(rois, 0, None)  # several coordinates are negative

        assert len(rois) == len(ids)

        rois[:, 2:] += rois[:, :2]
        wid += len(ids)
        nid += sum(ids == -2)
    print('wid: {}, nid: {}'.format(wid, nid))

    # extr_anno = "../prw_feats/extract_prw_anno_yolo.pkl" #get from yolo-prw
    dir = '../0yolov5_scores/'
    extr_anno = dir + cfg.score_anno
    print('loading prw yolov5l boxes')
    annotations = unpickle(extr_anno)
    all_boxes=0
    for j,a in enumerate(annotations):
        img_name = a['img_path'].split('/')[-1]
        img_path = osp.join(data_imgroot,img_name)
        annotations[j]['img_path']=img_path
        num_boxes =len(a['pids'])
        all_boxes+=num_boxes

    mem = torch.zeros([all_boxes, 256])
    print('loading zeros mem')
    mem=mem.to(cfg.DEVICE)

    num_features = 256
    memory = Memory(num_features, num_classes=[], mem=mem)
    memory.to(cfg.DEVICE)
    return annotations, memory

def create_new_pid_ssm_yolo5l(cfg):
    annotations = []
    img_has_boxes = {}  # use yolo detector {path:boxes}

    yolobox_json = cfg.score_anno

    with open(osp.join(yolobox_json), "r") as f:
        yolodet_boxes = json.load(f)  # {'person':[[x1,y1,x2,y2,s,img_path]]}
    num_classes = len(yolodet_boxes['person'])

    for box_info in yolodet_boxes['person']:
        box, img_path = box_info[:4], box_info[-1]
        img_has_boxes[img_path] = []
    for box_info in yolodet_boxes['person']:
        box, img_path = box_info[:4], box_info[-1]
        img_has_boxes[img_path].append(box)

    print('yolo总boxes数： ', num_classes)

    # 提取boxes-anno
    pid = 1
    for img_p, boxes in img_has_boxes.items():
        img_name = img_p.split('/')[-1]
        pids = [pid + x for x in range(len(boxes))]  # 3
        pid = pids[-1] + 1  # 4
        annotations.append(
            {
                "img_name": img_name,
                "img_path": osp.join(cfg.INPUT.DATA_ROOT, 'Image', 'SSM', img_name),
                "boxes": boxes,
                "pids": pids
            }
        )

    mem = torch.zeros([num_classes, 256])
    print('loading zeros mem finished..')
    mem = mem.to(cfg.DEVICE)

    num_features = 256
    memory = Memory(num_features, num_classes, mem)
    memory.to(cfg.DEVICE)

    return annotations, memory